{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set GPU clocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo nvidia-persistenced\n",
    "!sudo nvidia-smi -ac 877,1530"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from core import *\n",
    "from torch_backend import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_bn(c_in, c_out):\n",
    "    return {\n",
    "        'conv': nn.Conv2d(c_in, c_out, kernel_size=3, stride=1, padding=1, bias=False), \n",
    "        'bn': BatchNorm(c_out), \n",
    "        'relu': nn.ReLU(True)\n",
    "    }\n",
    "\n",
    "def residual(c):\n",
    "    return {\n",
    "        'in': Identity(),\n",
    "        'res1': conv_bn(c, c),\n",
    "        'res2': conv_bn(c, c),\n",
    "        'add': (Add(), ['in', 'res2/relu']),\n",
    "    }\n",
    "\n",
    "def net(channels=None, weight=0.125, pool=nn.MaxPool2d(2), extra_layers=(), res_layers=('layer1', 'layer3')):\n",
    "    channels = channels or {'prep': 64, 'layer1': 128, 'layer2': 256, 'layer3': 512}\n",
    "    n = {\n",
    "        'input': (None, []),\n",
    "        'prep': conv_bn(3, channels['prep']),\n",
    "        'layer1': dict(conv_bn(channels['prep'], channels['layer1']), pool=pool),\n",
    "        'layer2': dict(conv_bn(channels['layer1'], channels['layer2']), pool=pool),\n",
    "        'layer3': dict(conv_bn(channels['layer2'], channels['layer3']), pool=pool),\n",
    "        'pool': nn.MaxPool2d(4),\n",
    "        'flatten': Flatten(),\n",
    "        'linear': nn.Linear(channels['layer3'], 10, bias=False),\n",
    "        'logits': Mul(weight),\n",
    "    }\n",
    "    for layer in res_layers:\n",
    "        n[layer]['residual'] = residual(channels[layer])\n",
    "    for layer in extra_layers:\n",
    "        n[layer]['extra'] = conv_bn(channels[layer], channels[layer])       \n",
    "    return n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download and preprocess data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Preprocessing training data\n",
      "Finished in 2.9 seconds\n",
      "Preprocessing test data\n",
      "Finished in 0.15 seconds\n"
     ]
    }
   ],
   "source": [
    "DATA_DIR = './data'\n",
    "dataset = cifar10(root=DATA_DIR)\n",
    "timer = Timer()\n",
    "print('Preprocessing training data')\n",
    "transforms = [\n",
    "    partial(normalise, mean=np.array(cifar10_mean, dtype=np.float32), std=np.array(cifar10_std, dtype=np.float32)),\n",
    "    partial(transpose, source='NHWC', target='NCHW'), \n",
    "]\n",
    "train_set = list(zip(*preprocess(dataset['train'], [partial(pad, border=4)] + transforms).values()))\n",
    "print(f'Finished in {timer():.2} seconds')\n",
    "print('Preprocessing test data')\n",
    "valid_set = list(zip(*preprocess(dataset['valid'], transforms).values()))\n",
    "print(f'Finished in {timer():.2} seconds')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network visualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.38.0 (20140413.2041)\n",
       " -->\n",
       "<!-- Title: G Pages: 1 -->\n",
       "<svg width=\"1080pt\" height=\"65pt\"\n",
       " viewBox=\"0.00 0.00 1080.00 65.23\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(0.336239 0.336239) rotate(0) translate(4 190)\">\n",
       "<title>G</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-190 3208,-190 3208,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\"><title>cluster_prep</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M94,-44C94,-44 320,-44 320,-44 326,-44 332,-50 332,-56 332,-56 332,-106 332,-106 332,-112 326,-118 320,-118 320,-118 94,-118 94,-118 88,-118 82,-112 82,-106 82,-106 82,-56 82,-56 82,-50 88,-44 94,-44\"/>\n",
       "<text text-anchor=\"middle\" x=\"207\" y=\"-102.8\" font-family=\"Times,serif\" font-size=\"14.00\">prep</text>\n",
       "</g>\n",
       "<g id=\"clust2\" class=\"cluster\"><title>cluster_layer1</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M364,-8C364,-8 1408,-8 1408,-8 1414,-8 1420,-14 1420,-20 1420,-20 1420,-156 1420,-156 1420,-162 1414,-168 1408,-168 1408,-168 364,-168 364,-168 358,-168 352,-162 352,-156 352,-156 352,-20 352,-20 352,-14 358,-8 364,-8\"/>\n",
       "<text text-anchor=\"middle\" x=\"886\" y=\"-152.8\" font-family=\"Times,serif\" font-size=\"14.00\">layer1</text>\n",
       "</g>\n",
       "<g id=\"clust3\" class=\"cluster\"><title>cluster_layer1/residual</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M724,-16C724,-16 1400,-16 1400,-16 1406,-16 1412,-22 1412,-28 1412,-28 1412,-126 1412,-126 1412,-132 1406,-138 1400,-138 1400,-138 724,-138 724,-138 718,-138 712,-132 712,-126 712,-126 712,-28 712,-28 712,-22 718,-16 724,-16\"/>\n",
       "<text text-anchor=\"middle\" x=\"1062\" y=\"-122.8\" font-family=\"Times,serif\" font-size=\"14.00\">residual</text>\n",
       "</g>\n",
       "<g id=\"clust4\" class=\"cluster\"><title>cluster_layer1/residual/res1</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M814,-24C814,-24 1040,-24 1040,-24 1046,-24 1052,-30 1052,-36 1052,-36 1052,-86 1052,-86 1052,-92 1046,-98 1040,-98 1040,-98 814,-98 814,-98 808,-98 802,-92 802,-86 802,-86 802,-36 802,-36 802,-30 808,-24 814,-24\"/>\n",
       "<text text-anchor=\"middle\" x=\"927\" y=\"-82.8\" font-family=\"Times,serif\" font-size=\"14.00\">res1</text>\n",
       "</g>\n",
       "<g id=\"clust5\" class=\"cluster\"><title>cluster_layer1/residual/res2</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M1084,-24C1084,-24 1310,-24 1310,-24 1316,-24 1322,-30 1322,-36 1322,-36 1322,-86 1322,-86 1322,-92 1316,-98 1310,-98 1310,-98 1084,-98 1084,-98 1078,-98 1072,-92 1072,-86 1072,-86 1072,-36 1072,-36 1072,-30 1078,-24 1084,-24\"/>\n",
       "<text text-anchor=\"middle\" x=\"1197\" y=\"-82.8\" font-family=\"Times,serif\" font-size=\"14.00\">res2</text>\n",
       "</g>\n",
       "<g id=\"clust6\" class=\"cluster\"><title>cluster_layer2</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M1444,-44C1444,-44 1760,-44 1760,-44 1766,-44 1772,-50 1772,-56 1772,-56 1772,-106 1772,-106 1772,-112 1766,-118 1760,-118 1760,-118 1444,-118 1444,-118 1438,-118 1432,-112 1432,-106 1432,-106 1432,-56 1432,-56 1432,-50 1438,-44 1444,-44\"/>\n",
       "<text text-anchor=\"middle\" x=\"1602\" y=\"-102.8\" font-family=\"Times,serif\" font-size=\"14.00\">layer2</text>\n",
       "</g>\n",
       "<g id=\"clust7\" class=\"cluster\"><title>cluster_layer3</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M1804,-18C1804,-18 2848,-18 2848,-18 2854,-18 2860,-24 2860,-30 2860,-30 2860,-166 2860,-166 2860,-172 2854,-178 2848,-178 2848,-178 1804,-178 1804,-178 1798,-178 1792,-172 1792,-166 1792,-166 1792,-30 1792,-30 1792,-24 1798,-18 1804,-18\"/>\n",
       "<text text-anchor=\"middle\" x=\"2326\" y=\"-162.8\" font-family=\"Times,serif\" font-size=\"14.00\">layer3</text>\n",
       "</g>\n",
       "<g id=\"clust8\" class=\"cluster\"><title>cluster_layer3/residual</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M2164,-26C2164,-26 2840,-26 2840,-26 2846,-26 2852,-32 2852,-38 2852,-38 2852,-136 2852,-136 2852,-142 2846,-148 2840,-148 2840,-148 2164,-148 2164,-148 2158,-148 2152,-142 2152,-136 2152,-136 2152,-38 2152,-38 2152,-32 2158,-26 2164,-26\"/>\n",
       "<text text-anchor=\"middle\" x=\"2502\" y=\"-132.8\" font-family=\"Times,serif\" font-size=\"14.00\">residual</text>\n",
       "</g>\n",
       "<g id=\"clust9\" class=\"cluster\"><title>cluster_layer3/residual/res1</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M2254,-34C2254,-34 2480,-34 2480,-34 2486,-34 2492,-40 2492,-46 2492,-46 2492,-96 2492,-96 2492,-102 2486,-108 2480,-108 2480,-108 2254,-108 2254,-108 2248,-108 2242,-102 2242,-96 2242,-96 2242,-46 2242,-46 2242,-40 2248,-34 2254,-34\"/>\n",
       "<text text-anchor=\"middle\" x=\"2367\" y=\"-92.8\" font-family=\"Times,serif\" font-size=\"14.00\">res1</text>\n",
       "</g>\n",
       "<g id=\"clust10\" class=\"cluster\"><title>cluster_layer3/residual/res2</title>\n",
       "<path fill=\"#777777\" fill-opacity=\"0.266667\" stroke=\"black\" d=\"M2524,-34C2524,-34 2750,-34 2750,-34 2756,-34 2762,-40 2762,-46 2762,-46 2762,-96 2762,-96 2762,-102 2756,-108 2750,-108 2750,-108 2524,-108 2524,-108 2518,-108 2512,-102 2512,-96 2512,-96 2512,-46 2512,-46 2512,-40 2518,-34 2524,-34\"/>\n",
       "<text text-anchor=\"middle\" x=\"2637\" y=\"-92.8\" font-family=\"Times,serif\" font-size=\"14.00\">res2</text>\n",
       "</g>\n",
       "<!-- prep/conv -->\n",
       "<g id=\"node1\" class=\"node\"><title>prep/conv</title>\n",
       "<g id=\"a_node1\"><a xlink:title=\"Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M132,-88C132,-88 102,-88 102,-88 96,-88 90,-82 90,-76 90,-76 90,-64 90,-64 90,-58 96,-52 102,-52 102,-52 132,-52 132,-52 138,-52 144,-58 144,-64 144,-64 144,-76 144,-76 144,-82 138,-88 132,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"117\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- prep/bn -->\n",
       "<g id=\"node2\" class=\"node\"><title>prep/bn</title>\n",
       "<g id=\"a_node2\"><a xlink:title=\"BatchNorm(64, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M222,-88C222,-88 192,-88 192,-88 186,-88 180,-82 180,-76 180,-76 180,-64 180,-64 180,-58 186,-52 192,-52 192,-52 222,-52 222,-52 228,-52 234,-58 234,-64 234,-64 234,-76 234,-76 234,-82 228,-88 222,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"207\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- prep/conv&#45;&gt;prep/bn -->\n",
       "<g id=\"edge2\" class=\"edge\"><title>prep/conv&#45;&gt;prep/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M144.403,-70C152.393,-70 161.311,-70 169.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"169.919,-73.5001 179.919,-70 169.919,-66.5001 169.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- prep/relu -->\n",
       "<g id=\"node3\" class=\"node\"><title>prep/relu</title>\n",
       "<g id=\"a_node3\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M312,-88C312,-88 282,-88 282,-88 276,-88 270,-82 270,-76 270,-76 270,-64 270,-64 270,-58 276,-52 282,-52 282,-52 312,-52 312,-52 318,-52 324,-58 324,-64 324,-64 324,-76 324,-76 324,-82 318,-88 312,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"297\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- prep/bn&#45;&gt;prep/relu -->\n",
       "<g id=\"edge3\" class=\"edge\"><title>prep/bn&#45;&gt;prep/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M234.403,-70C242.393,-70 251.311,-70 259.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"259.919,-73.5001 269.919,-70 259.919,-66.5001 259.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/conv -->\n",
       "<g id=\"node4\" class=\"node\"><title>layer1/conv</title>\n",
       "<g id=\"a_node4\"><a xlink:title=\"Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M402,-88C402,-88 372,-88 372,-88 366,-88 360,-82 360,-76 360,-76 360,-64 360,-64 360,-58 366,-52 372,-52 372,-52 402,-52 402,-52 408,-52 414,-58 414,-64 414,-64 414,-76 414,-76 414,-82 408,-88 402,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"387\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- prep/relu&#45;&gt;layer1/conv -->\n",
       "<g id=\"edge4\" class=\"edge\"><title>prep/relu&#45;&gt;layer1/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M324.403,-70C332.393,-70 341.311,-70 349.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"349.919,-73.5001 359.919,-70 349.919,-66.5001 349.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/bn -->\n",
       "<g id=\"node5\" class=\"node\"><title>layer1/bn</title>\n",
       "<g id=\"a_node5\"><a xlink:title=\"BatchNorm(128, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M492,-88C492,-88 462,-88 462,-88 456,-88 450,-82 450,-76 450,-76 450,-64 450,-64 450,-58 456,-52 462,-52 462,-52 492,-52 492,-52 498,-52 504,-58 504,-64 504,-64 504,-76 504,-76 504,-82 498,-88 492,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"477\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/conv&#45;&gt;layer1/bn -->\n",
       "<g id=\"edge5\" class=\"edge\"><title>layer1/conv&#45;&gt;layer1/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M414.403,-70C422.393,-70 431.311,-70 439.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"439.919,-73.5001 449.919,-70 439.919,-66.5001 439.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/relu -->\n",
       "<g id=\"node6\" class=\"node\"><title>layer1/relu</title>\n",
       "<g id=\"a_node6\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M582,-88C582,-88 552,-88 552,-88 546,-88 540,-82 540,-76 540,-76 540,-64 540,-64 540,-58 546,-52 552,-52 552,-52 582,-52 582,-52 588,-52 594,-58 594,-64 594,-64 594,-76 594,-76 594,-82 588,-88 582,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"567\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/bn&#45;&gt;layer1/relu -->\n",
       "<g id=\"edge6\" class=\"edge\"><title>layer1/bn&#45;&gt;layer1/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M504.403,-70C512.393,-70 521.311,-70 529.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"529.919,-73.5001 539.919,-70 529.919,-66.5001 529.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/pool -->\n",
       "<g id=\"node7\" class=\"node\"><title>layer1/pool</title>\n",
       "<g id=\"a_node7\"><a xlink:title=\"MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\">\n",
       "<path fill=\"#8dd3c7\" stroke=\"black\" d=\"M672,-88C672,-88 642,-88 642,-88 636,-88 630,-82 630,-76 630,-76 630,-64 630,-64 630,-58 636,-52 642,-52 642,-52 672,-52 672,-52 678,-52 684,-58 684,-64 684,-64 684,-76 684,-76 684,-82 678,-88 672,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"657\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">pool</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/relu&#45;&gt;layer1/pool -->\n",
       "<g id=\"edge7\" class=\"edge\"><title>layer1/relu&#45;&gt;layer1/pool</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M594.403,-70C602.393,-70 611.311,-70 619.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"619.919,-73.5001 629.919,-70 619.919,-66.5001 619.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/in -->\n",
       "<g id=\"node8\" class=\"node\"><title>layer1/residual/in</title>\n",
       "<g id=\"a_node8\"><a xlink:title=\"Identity()\">\n",
       "<path fill=\"#80b1d3\" stroke=\"black\" d=\"M762,-88C762,-88 732,-88 732,-88 726,-88 720,-82 720,-76 720,-76 720,-64 720,-64 720,-58 726,-52 732,-52 732,-52 762,-52 762,-52 768,-52 774,-58 774,-64 774,-64 774,-76 774,-76 774,-82 768,-88 762,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"747\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">in</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/pool&#45;&gt;layer1/residual/in -->\n",
       "<g id=\"edge8\" class=\"edge\"><title>layer1/pool&#45;&gt;layer1/residual/in</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M684.403,-70C692.393,-70 701.311,-70 709.824,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"709.919,-73.5001 719.919,-70 709.919,-66.5001 709.919,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res1/conv -->\n",
       "<g id=\"node9\" class=\"node\"><title>layer1/residual/res1/conv</title>\n",
       "<g id=\"a_node9\"><a xlink:title=\"Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M852,-68C852,-68 822,-68 822,-68 816,-68 810,-62 810,-56 810,-56 810,-44 810,-44 810,-38 816,-32 822,-32 822,-32 852,-32 852,-32 858,-32 864,-38 864,-44 864,-44 864,-56 864,-56 864,-62 858,-68 852,-68\"/>\n",
       "<text text-anchor=\"middle\" x=\"837\" y=\"-45.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/in&#45;&gt;layer1/residual/res1/conv -->\n",
       "<g id=\"edge9\" class=\"edge\"><title>layer1/residual/in&#45;&gt;layer1/residual/res1/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M774.403,-63.9993C782.481,-62.1634 791.507,-60.1121 800.105,-58.158\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"800.943,-61.5568 809.919,-55.9275 799.392,-54.7308 800.943,-61.5568\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/add -->\n",
       "<g id=\"node15\" class=\"node\"><title>layer1/residual/add</title>\n",
       "<g id=\"a_node15\"><a xlink:title=\"Add()\">\n",
       "<path fill=\"#fdb462\" stroke=\"black\" d=\"M1392,-88C1392,-88 1362,-88 1362,-88 1356,-88 1350,-82 1350,-76 1350,-76 1350,-64 1350,-64 1350,-58 1356,-52 1362,-52 1362,-52 1392,-52 1392,-52 1398,-52 1404,-58 1404,-64 1404,-64 1404,-76 1404,-76 1404,-82 1398,-88 1392,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1377\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">add</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/in&#45;&gt;layer1/residual/add -->\n",
       "<g id=\"edge15\" class=\"edge\"><title>layer1/residual/in&#45;&gt;layer1/residual/add</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M772.616,-88.06C781.446,-93.5709 791.759,-98.9705 802,-102 854.89,-117.647 870.844,-107 926,-107 926,-107 926,-107 1198,-107 1253.16,-107 1269.11,-117.647 1322,-102 1329.04,-99.9172 1336.12,-96.7142 1342.72,-93.1249\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1344.52,-96.128 1351.38,-88.06 1340.98,-90.0846 1344.52,-96.128\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res1/bn -->\n",
       "<g id=\"node10\" class=\"node\"><title>layer1/residual/res1/bn</title>\n",
       "<g id=\"a_node10\"><a xlink:title=\"BatchNorm(128, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M942,-68C942,-68 912,-68 912,-68 906,-68 900,-62 900,-56 900,-56 900,-44 900,-44 900,-38 906,-32 912,-32 912,-32 942,-32 942,-32 948,-32 954,-38 954,-44 954,-44 954,-56 954,-56 954,-62 948,-68 942,-68\"/>\n",
       "<text text-anchor=\"middle\" x=\"927\" y=\"-45.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/res1/conv&#45;&gt;layer1/residual/res1/bn -->\n",
       "<g id=\"edge10\" class=\"edge\"><title>layer1/residual/res1/conv&#45;&gt;layer1/residual/res1/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M864.403,-50C872.393,-50 881.311,-50 889.824,-50\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"889.919,-53.5001 899.919,-50 889.919,-46.5001 889.919,-53.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res1/relu -->\n",
       "<g id=\"node11\" class=\"node\"><title>layer1/residual/res1/relu</title>\n",
       "<g id=\"a_node11\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M1032,-68C1032,-68 1002,-68 1002,-68 996,-68 990,-62 990,-56 990,-56 990,-44 990,-44 990,-38 996,-32 1002,-32 1002,-32 1032,-32 1032,-32 1038,-32 1044,-38 1044,-44 1044,-44 1044,-56 1044,-56 1044,-62 1038,-68 1032,-68\"/>\n",
       "<text text-anchor=\"middle\" x=\"1017\" y=\"-45.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/res1/bn&#45;&gt;layer1/residual/res1/relu -->\n",
       "<g id=\"edge11\" class=\"edge\"><title>layer1/residual/res1/bn&#45;&gt;layer1/residual/res1/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M954.403,-50C962.393,-50 971.311,-50 979.824,-50\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"979.919,-53.5001 989.919,-50 979.919,-46.5001 979.919,-53.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res2/conv -->\n",
       "<g id=\"node12\" class=\"node\"><title>layer1/residual/res2/conv</title>\n",
       "<g id=\"a_node12\"><a xlink:title=\"Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M1122,-68C1122,-68 1092,-68 1092,-68 1086,-68 1080,-62 1080,-56 1080,-56 1080,-44 1080,-44 1080,-38 1086,-32 1092,-32 1092,-32 1122,-32 1122,-32 1128,-32 1134,-38 1134,-44 1134,-44 1134,-56 1134,-56 1134,-62 1128,-68 1122,-68\"/>\n",
       "<text text-anchor=\"middle\" x=\"1107\" y=\"-45.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/res1/relu&#45;&gt;layer1/residual/res2/conv -->\n",
       "<g id=\"edge12\" class=\"edge\"><title>layer1/residual/res1/relu&#45;&gt;layer1/residual/res2/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1044.4,-50C1052.39,-50 1061.31,-50 1069.82,-50\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1069.92,-53.5001 1079.92,-50 1069.92,-46.5001 1069.92,-53.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res2/bn -->\n",
       "<g id=\"node13\" class=\"node\"><title>layer1/residual/res2/bn</title>\n",
       "<g id=\"a_node13\"><a xlink:title=\"BatchNorm(128, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M1212,-68C1212,-68 1182,-68 1182,-68 1176,-68 1170,-62 1170,-56 1170,-56 1170,-44 1170,-44 1170,-38 1176,-32 1182,-32 1182,-32 1212,-32 1212,-32 1218,-32 1224,-38 1224,-44 1224,-44 1224,-56 1224,-56 1224,-62 1218,-68 1212,-68\"/>\n",
       "<text text-anchor=\"middle\" x=\"1197\" y=\"-45.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/res2/conv&#45;&gt;layer1/residual/res2/bn -->\n",
       "<g id=\"edge13\" class=\"edge\"><title>layer1/residual/res2/conv&#45;&gt;layer1/residual/res2/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1134.4,-50C1142.39,-50 1151.31,-50 1159.82,-50\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1159.92,-53.5001 1169.92,-50 1159.92,-46.5001 1159.92,-53.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res2/relu -->\n",
       "<g id=\"node14\" class=\"node\"><title>layer1/residual/res2/relu</title>\n",
       "<g id=\"a_node14\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M1302,-68C1302,-68 1272,-68 1272,-68 1266,-68 1260,-62 1260,-56 1260,-56 1260,-44 1260,-44 1260,-38 1266,-32 1272,-32 1272,-32 1302,-32 1302,-32 1308,-32 1314,-38 1314,-44 1314,-44 1314,-56 1314,-56 1314,-62 1308,-68 1302,-68\"/>\n",
       "<text text-anchor=\"middle\" x=\"1287\" y=\"-45.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/res2/bn&#45;&gt;layer1/residual/res2/relu -->\n",
       "<g id=\"edge14\" class=\"edge\"><title>layer1/residual/res2/bn&#45;&gt;layer1/residual/res2/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1224.4,-50C1232.39,-50 1241.31,-50 1249.82,-50\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1249.92,-53.5001 1259.92,-50 1249.92,-46.5001 1249.92,-53.5001\"/>\n",
       "</g>\n",
       "<!-- layer1/residual/res2/relu&#45;&gt;layer1/residual/add -->\n",
       "<g id=\"edge16\" class=\"edge\"><title>layer1/residual/res2/relu&#45;&gt;layer1/residual/add</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1314.4,-56.0007C1322.48,-57.8366 1331.51,-59.8879 1340.1,-61.842\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1339.39,-65.2692 1349.92,-64.0725 1340.94,-58.4432 1339.39,-65.2692\"/>\n",
       "</g>\n",
       "<!-- layer2/conv -->\n",
       "<g id=\"node16\" class=\"node\"><title>layer2/conv</title>\n",
       "<g id=\"a_node16\"><a xlink:title=\"Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M1482,-88C1482,-88 1452,-88 1452,-88 1446,-88 1440,-82 1440,-76 1440,-76 1440,-64 1440,-64 1440,-58 1446,-52 1452,-52 1452,-52 1482,-52 1482,-52 1488,-52 1494,-58 1494,-64 1494,-64 1494,-76 1494,-76 1494,-82 1488,-88 1482,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1467\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer1/residual/add&#45;&gt;layer2/conv -->\n",
       "<g id=\"edge17\" class=\"edge\"><title>layer1/residual/add&#45;&gt;layer2/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1404.4,-70C1412.39,-70 1421.31,-70 1429.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1429.92,-73.5001 1439.92,-70 1429.92,-66.5001 1429.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer2/bn -->\n",
       "<g id=\"node17\" class=\"node\"><title>layer2/bn</title>\n",
       "<g id=\"a_node17\"><a xlink:title=\"BatchNorm(256, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M1572,-88C1572,-88 1542,-88 1542,-88 1536,-88 1530,-82 1530,-76 1530,-76 1530,-64 1530,-64 1530,-58 1536,-52 1542,-52 1542,-52 1572,-52 1572,-52 1578,-52 1584,-58 1584,-64 1584,-64 1584,-76 1584,-76 1584,-82 1578,-88 1572,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1557\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer2/conv&#45;&gt;layer2/bn -->\n",
       "<g id=\"edge18\" class=\"edge\"><title>layer2/conv&#45;&gt;layer2/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1494.4,-70C1502.39,-70 1511.31,-70 1519.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1519.92,-73.5001 1529.92,-70 1519.92,-66.5001 1519.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer2/relu -->\n",
       "<g id=\"node18\" class=\"node\"><title>layer2/relu</title>\n",
       "<g id=\"a_node18\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M1662,-88C1662,-88 1632,-88 1632,-88 1626,-88 1620,-82 1620,-76 1620,-76 1620,-64 1620,-64 1620,-58 1626,-52 1632,-52 1632,-52 1662,-52 1662,-52 1668,-52 1674,-58 1674,-64 1674,-64 1674,-76 1674,-76 1674,-82 1668,-88 1662,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1647\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer2/bn&#45;&gt;layer2/relu -->\n",
       "<g id=\"edge19\" class=\"edge\"><title>layer2/bn&#45;&gt;layer2/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1584.4,-70C1592.39,-70 1601.31,-70 1609.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1609.92,-73.5001 1619.92,-70 1609.92,-66.5001 1609.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer2/pool -->\n",
       "<g id=\"node19\" class=\"node\"><title>layer2/pool</title>\n",
       "<g id=\"a_node19\"><a xlink:title=\"MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\">\n",
       "<path fill=\"#8dd3c7\" stroke=\"black\" d=\"M1752,-88C1752,-88 1722,-88 1722,-88 1716,-88 1710,-82 1710,-76 1710,-76 1710,-64 1710,-64 1710,-58 1716,-52 1722,-52 1722,-52 1752,-52 1752,-52 1758,-52 1764,-58 1764,-64 1764,-64 1764,-76 1764,-76 1764,-82 1758,-88 1752,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1737\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">pool</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer2/relu&#45;&gt;layer2/pool -->\n",
       "<g id=\"edge20\" class=\"edge\"><title>layer2/relu&#45;&gt;layer2/pool</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1674.4,-70C1682.39,-70 1691.31,-70 1699.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1699.92,-73.5001 1709.92,-70 1699.92,-66.5001 1699.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/conv -->\n",
       "<g id=\"node20\" class=\"node\"><title>layer3/conv</title>\n",
       "<g id=\"a_node20\"><a xlink:title=\"Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M1842,-88C1842,-88 1812,-88 1812,-88 1806,-88 1800,-82 1800,-76 1800,-76 1800,-64 1800,-64 1800,-58 1806,-52 1812,-52 1812,-52 1842,-52 1842,-52 1848,-52 1854,-58 1854,-64 1854,-64 1854,-76 1854,-76 1854,-82 1848,-88 1842,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1827\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer2/pool&#45;&gt;layer3/conv -->\n",
       "<g id=\"edge21\" class=\"edge\"><title>layer2/pool&#45;&gt;layer3/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1764.4,-70C1772.39,-70 1781.31,-70 1789.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1789.92,-73.5001 1799.92,-70 1789.92,-66.5001 1789.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/bn -->\n",
       "<g id=\"node21\" class=\"node\"><title>layer3/bn</title>\n",
       "<g id=\"a_node21\"><a xlink:title=\"BatchNorm(512, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M1932,-88C1932,-88 1902,-88 1902,-88 1896,-88 1890,-82 1890,-76 1890,-76 1890,-64 1890,-64 1890,-58 1896,-52 1902,-52 1902,-52 1932,-52 1932,-52 1938,-52 1944,-58 1944,-64 1944,-64 1944,-76 1944,-76 1944,-82 1938,-88 1932,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"1917\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/conv&#45;&gt;layer3/bn -->\n",
       "<g id=\"edge22\" class=\"edge\"><title>layer3/conv&#45;&gt;layer3/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1854.4,-70C1862.39,-70 1871.31,-70 1879.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1879.92,-73.5001 1889.92,-70 1879.92,-66.5001 1879.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/relu -->\n",
       "<g id=\"node22\" class=\"node\"><title>layer3/relu</title>\n",
       "<g id=\"a_node22\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M2022,-88C2022,-88 1992,-88 1992,-88 1986,-88 1980,-82 1980,-76 1980,-76 1980,-64 1980,-64 1980,-58 1986,-52 1992,-52 1992,-52 2022,-52 2022,-52 2028,-52 2034,-58 2034,-64 2034,-64 2034,-76 2034,-76 2034,-82 2028,-88 2022,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"2007\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/bn&#45;&gt;layer3/relu -->\n",
       "<g id=\"edge23\" class=\"edge\"><title>layer3/bn&#45;&gt;layer3/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1944.4,-70C1952.39,-70 1961.31,-70 1969.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1969.92,-73.5001 1979.92,-70 1969.92,-66.5001 1969.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/pool -->\n",
       "<g id=\"node23\" class=\"node\"><title>layer3/pool</title>\n",
       "<g id=\"a_node23\"><a xlink:title=\"MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\">\n",
       "<path fill=\"#8dd3c7\" stroke=\"black\" d=\"M2112,-88C2112,-88 2082,-88 2082,-88 2076,-88 2070,-82 2070,-76 2070,-76 2070,-64 2070,-64 2070,-58 2076,-52 2082,-52 2082,-52 2112,-52 2112,-52 2118,-52 2124,-58 2124,-64 2124,-64 2124,-76 2124,-76 2124,-82 2118,-88 2112,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"2097\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">pool</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/relu&#45;&gt;layer3/pool -->\n",
       "<g id=\"edge24\" class=\"edge\"><title>layer3/relu&#45;&gt;layer3/pool</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2034.4,-70C2042.39,-70 2051.31,-70 2059.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2059.92,-73.5001 2069.92,-70 2059.92,-66.5001 2059.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/in -->\n",
       "<g id=\"node24\" class=\"node\"><title>layer3/residual/in</title>\n",
       "<g id=\"a_node24\"><a xlink:title=\"Identity()\">\n",
       "<path fill=\"#80b1d3\" stroke=\"black\" d=\"M2202,-88C2202,-88 2172,-88 2172,-88 2166,-88 2160,-82 2160,-76 2160,-76 2160,-64 2160,-64 2160,-58 2166,-52 2172,-52 2172,-52 2202,-52 2202,-52 2208,-52 2214,-58 2214,-64 2214,-64 2214,-76 2214,-76 2214,-82 2208,-88 2202,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"2187\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">in</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/pool&#45;&gt;layer3/residual/in -->\n",
       "<g id=\"edge25\" class=\"edge\"><title>layer3/pool&#45;&gt;layer3/residual/in</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2124.4,-70C2132.39,-70 2141.31,-70 2149.82,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2149.92,-73.5001 2159.92,-70 2149.92,-66.5001 2149.92,-73.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res1/conv -->\n",
       "<g id=\"node25\" class=\"node\"><title>layer3/residual/res1/conv</title>\n",
       "<g id=\"a_node25\"><a xlink:title=\"Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M2292,-78C2292,-78 2262,-78 2262,-78 2256,-78 2250,-72 2250,-66 2250,-66 2250,-54 2250,-54 2250,-48 2256,-42 2262,-42 2262,-42 2292,-42 2292,-42 2298,-42 2304,-48 2304,-54 2304,-54 2304,-66 2304,-66 2304,-72 2298,-78 2292,-78\"/>\n",
       "<text text-anchor=\"middle\" x=\"2277\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/in&#45;&gt;layer3/residual/res1/conv -->\n",
       "<g id=\"edge26\" class=\"edge\"><title>layer3/residual/in&#45;&gt;layer3/residual/res1/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2214.4,-66.9997C2222.39,-66.0917 2231.31,-65.0783 2239.82,-64.1109\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2240.38,-67.5706 2249.92,-62.9638 2239.59,-60.6153 2240.38,-67.5706\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/add -->\n",
       "<g id=\"node31\" class=\"node\"><title>layer3/residual/add</title>\n",
       "<g id=\"a_node31\"><a xlink:title=\"Add()\">\n",
       "<path fill=\"#fdb462\" stroke=\"black\" d=\"M2832,-98C2832,-98 2802,-98 2802,-98 2796,-98 2790,-92 2790,-86 2790,-86 2790,-74 2790,-74 2790,-68 2796,-62 2802,-62 2802,-62 2832,-62 2832,-62 2838,-62 2844,-68 2844,-74 2844,-74 2844,-86 2844,-86 2844,-92 2838,-98 2832,-98\"/>\n",
       "<text text-anchor=\"middle\" x=\"2817\" y=\"-75.8\" font-family=\"Times,serif\" font-size=\"14.00\">add</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/in&#45;&gt;layer3/residual/add -->\n",
       "<g id=\"edge32\" class=\"edge\"><title>layer3/residual/in&#45;&gt;layer3/residual/add</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2205.71,-88.3841C2215.49,-97.2503 2228.39,-107.022 2242,-112 2293.8,-130.948 2310.84,-117 2366,-117 2366,-117 2366,-117 2638,-117 2693.16,-117 2709.11,-127.647 2762,-112 2769.04,-109.917 2776.12,-106.714 2782.72,-103.125\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2784.52,-106.128 2791.38,-98.06 2780.98,-100.085 2784.52,-106.128\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res1/bn -->\n",
       "<g id=\"node26\" class=\"node\"><title>layer3/residual/res1/bn</title>\n",
       "<g id=\"a_node26\"><a xlink:title=\"BatchNorm(512, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M2382,-78C2382,-78 2352,-78 2352,-78 2346,-78 2340,-72 2340,-66 2340,-66 2340,-54 2340,-54 2340,-48 2346,-42 2352,-42 2352,-42 2382,-42 2382,-42 2388,-42 2394,-48 2394,-54 2394,-54 2394,-66 2394,-66 2394,-72 2388,-78 2382,-78\"/>\n",
       "<text text-anchor=\"middle\" x=\"2367\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/res1/conv&#45;&gt;layer3/residual/res1/bn -->\n",
       "<g id=\"edge27\" class=\"edge\"><title>layer3/residual/res1/conv&#45;&gt;layer3/residual/res1/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2304.4,-60C2312.39,-60 2321.31,-60 2329.82,-60\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2329.92,-63.5001 2339.92,-60 2329.92,-56.5001 2329.92,-63.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res1/relu -->\n",
       "<g id=\"node27\" class=\"node\"><title>layer3/residual/res1/relu</title>\n",
       "<g id=\"a_node27\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M2472,-78C2472,-78 2442,-78 2442,-78 2436,-78 2430,-72 2430,-66 2430,-66 2430,-54 2430,-54 2430,-48 2436,-42 2442,-42 2442,-42 2472,-42 2472,-42 2478,-42 2484,-48 2484,-54 2484,-54 2484,-66 2484,-66 2484,-72 2478,-78 2472,-78\"/>\n",
       "<text text-anchor=\"middle\" x=\"2457\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/res1/bn&#45;&gt;layer3/residual/res1/relu -->\n",
       "<g id=\"edge28\" class=\"edge\"><title>layer3/residual/res1/bn&#45;&gt;layer3/residual/res1/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2394.4,-60C2402.39,-60 2411.31,-60 2419.82,-60\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2419.92,-63.5001 2429.92,-60 2419.92,-56.5001 2419.92,-63.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res2/conv -->\n",
       "<g id=\"node28\" class=\"node\"><title>layer3/residual/res2/conv</title>\n",
       "<g id=\"a_node28\"><a xlink:title=\"Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\">\n",
       "<path fill=\"#bebada\" stroke=\"black\" d=\"M2562,-78C2562,-78 2532,-78 2532,-78 2526,-78 2520,-72 2520,-66 2520,-66 2520,-54 2520,-54 2520,-48 2526,-42 2532,-42 2532,-42 2562,-42 2562,-42 2568,-42 2574,-48 2574,-54 2574,-54 2574,-66 2574,-66 2574,-72 2568,-78 2562,-78\"/>\n",
       "<text text-anchor=\"middle\" x=\"2547\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">conv</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/res1/relu&#45;&gt;layer3/residual/res2/conv -->\n",
       "<g id=\"edge29\" class=\"edge\"><title>layer3/residual/res1/relu&#45;&gt;layer3/residual/res2/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2484.4,-60C2492.39,-60 2501.31,-60 2509.82,-60\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2509.92,-63.5001 2519.92,-60 2509.92,-56.5001 2509.92,-63.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res2/bn -->\n",
       "<g id=\"node29\" class=\"node\"><title>layer3/residual/res2/bn</title>\n",
       "<g id=\"a_node29\"><a xlink:title=\"BatchNorm(512, eps=1e&#45;05, momentum=0.1, affine=True, track_running_stats=True)\">\n",
       "<path fill=\"#ffffb3\" stroke=\"black\" d=\"M2652,-78C2652,-78 2622,-78 2622,-78 2616,-78 2610,-72 2610,-66 2610,-66 2610,-54 2610,-54 2610,-48 2616,-42 2622,-42 2622,-42 2652,-42 2652,-42 2658,-42 2664,-48 2664,-54 2664,-54 2664,-66 2664,-66 2664,-72 2658,-78 2652,-78\"/>\n",
       "<text text-anchor=\"middle\" x=\"2637\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">bn</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/res2/conv&#45;&gt;layer3/residual/res2/bn -->\n",
       "<g id=\"edge30\" class=\"edge\"><title>layer3/residual/res2/conv&#45;&gt;layer3/residual/res2/bn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2574.4,-60C2582.39,-60 2591.31,-60 2599.82,-60\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2599.92,-63.5001 2609.92,-60 2599.92,-56.5001 2599.92,-63.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res2/relu -->\n",
       "<g id=\"node30\" class=\"node\"><title>layer3/residual/res2/relu</title>\n",
       "<g id=\"a_node30\"><a xlink:title=\"ReLU(inplace)\">\n",
       "<path fill=\"#fb8072\" stroke=\"black\" d=\"M2742,-78C2742,-78 2712,-78 2712,-78 2706,-78 2700,-72 2700,-66 2700,-66 2700,-54 2700,-54 2700,-48 2706,-42 2712,-42 2712,-42 2742,-42 2742,-42 2748,-42 2754,-48 2754,-54 2754,-54 2754,-66 2754,-66 2754,-72 2748,-78 2742,-78\"/>\n",
       "<text text-anchor=\"middle\" x=\"2727\" y=\"-55.8\" font-family=\"Times,serif\" font-size=\"14.00\">relu</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/res2/bn&#45;&gt;layer3/residual/res2/relu -->\n",
       "<g id=\"edge31\" class=\"edge\"><title>layer3/residual/res2/bn&#45;&gt;layer3/residual/res2/relu</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2664.4,-60C2672.39,-60 2681.31,-60 2689.82,-60\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2689.92,-63.5001 2699.92,-60 2689.92,-56.5001 2689.92,-63.5001\"/>\n",
       "</g>\n",
       "<!-- layer3/residual/res2/relu&#45;&gt;layer3/residual/add -->\n",
       "<g id=\"edge33\" class=\"edge\"><title>layer3/residual/res2/relu&#45;&gt;layer3/residual/add</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2754.4,-66.0007C2762.48,-67.8366 2771.51,-69.8879 2780.1,-71.842\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2779.39,-75.2692 2789.92,-74.0725 2780.94,-68.4432 2779.39,-75.2692\"/>\n",
       "</g>\n",
       "<!-- pool -->\n",
       "<g id=\"node32\" class=\"node\"><title>pool</title>\n",
       "<g id=\"a_node32\"><a xlink:title=\"MaxPool2d(kernel_size=4, stride=4, padding=0, dilation=1, ceil_mode=False)\">\n",
       "<path fill=\"#8dd3c7\" stroke=\"black\" d=\"M2922,-98C2922,-98 2892,-98 2892,-98 2886,-98 2880,-92 2880,-86 2880,-86 2880,-74 2880,-74 2880,-68 2886,-62 2892,-62 2892,-62 2922,-62 2922,-62 2928,-62 2934,-68 2934,-74 2934,-74 2934,-86 2934,-86 2934,-92 2928,-98 2922,-98\"/>\n",
       "<text text-anchor=\"middle\" x=\"2907\" y=\"-75.8\" font-family=\"Times,serif\" font-size=\"14.00\">pool</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- layer3/residual/add&#45;&gt;pool -->\n",
       "<g id=\"edge34\" class=\"edge\"><title>layer3/residual/add&#45;&gt;pool</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2844.4,-80C2852.39,-80 2861.31,-80 2869.82,-80\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2869.92,-83.5001 2879.92,-80 2869.92,-76.5001 2869.92,-83.5001\"/>\n",
       "</g>\n",
       "<!-- flatten -->\n",
       "<g id=\"node33\" class=\"node\"><title>flatten</title>\n",
       "<g id=\"a_node33\"><a xlink:title=\"Flatten()\">\n",
       "<path fill=\"#b3de69\" stroke=\"black\" d=\"M3012,-98C3012,-98 2982,-98 2982,-98 2976,-98 2970,-92 2970,-86 2970,-86 2970,-74 2970,-74 2970,-68 2976,-62 2982,-62 2982,-62 3012,-62 3012,-62 3018,-62 3024,-68 3024,-74 3024,-74 3024,-86 3024,-86 3024,-92 3018,-98 3012,-98\"/>\n",
       "<text text-anchor=\"middle\" x=\"2997\" y=\"-75.8\" font-family=\"Times,serif\" font-size=\"14.00\">flatten</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- pool&#45;&gt;flatten -->\n",
       "<g id=\"edge35\" class=\"edge\"><title>pool&#45;&gt;flatten</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M2934.4,-80C2942.39,-80 2951.31,-80 2959.82,-80\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"2959.92,-83.5001 2969.92,-80 2959.92,-76.5001 2959.92,-83.5001\"/>\n",
       "</g>\n",
       "<!-- linear -->\n",
       "<g id=\"node34\" class=\"node\"><title>linear</title>\n",
       "<g id=\"a_node34\"><a xlink:title=\"Linear(in_features=512, out_features=10, bias=False)\">\n",
       "<path fill=\"#fccde5\" stroke=\"black\" d=\"M3102,-98C3102,-98 3072,-98 3072,-98 3066,-98 3060,-92 3060,-86 3060,-86 3060,-74 3060,-74 3060,-68 3066,-62 3072,-62 3072,-62 3102,-62 3102,-62 3108,-62 3114,-68 3114,-74 3114,-74 3114,-86 3114,-86 3114,-92 3108,-98 3102,-98\"/>\n",
       "<text text-anchor=\"middle\" x=\"3087\" y=\"-75.8\" font-family=\"Times,serif\" font-size=\"14.00\">linear</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- flatten&#45;&gt;linear -->\n",
       "<g id=\"edge36\" class=\"edge\"><title>flatten&#45;&gt;linear</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M3024.4,-80C3032.39,-80 3041.31,-80 3049.82,-80\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"3049.92,-83.5001 3059.92,-80 3049.92,-76.5001 3049.92,-83.5001\"/>\n",
       "</g>\n",
       "<!-- logits -->\n",
       "<g id=\"node35\" class=\"node\"><title>logits</title>\n",
       "<g id=\"a_node35\"><a xlink:title=\"Mul()\">\n",
       "<path fill=\"#bc80bd\" stroke=\"black\" d=\"M3192,-98C3192,-98 3162,-98 3162,-98 3156,-98 3150,-92 3150,-86 3150,-86 3150,-74 3150,-74 3150,-68 3156,-62 3162,-62 3162,-62 3192,-62 3192,-62 3198,-62 3204,-68 3204,-74 3204,-74 3204,-86 3204,-86 3204,-92 3198,-98 3192,-98\"/>\n",
       "<text text-anchor=\"middle\" x=\"3177\" y=\"-75.8\" font-family=\"Times,serif\" font-size=\"14.00\">logits</text>\n",
       "</a>\n",
       "</g>\n",
       "</g>\n",
       "<!-- linear&#45;&gt;logits -->\n",
       "<g id=\"edge37\" class=\"edge\"><title>linear&#45;&gt;logits</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M3114.4,-80C3122.39,-80 3131.31,-80 3139.82,-80\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"3139.92,-83.5001 3149.92,-80 3139.92,-76.5001 3139.92,-83.5001\"/>\n",
       "</g>\n",
       "<!-- input -->\n",
       "<g id=\"node36\" class=\"node\"><title>input</title>\n",
       "<path fill=\"#ffffff\" stroke=\"black\" d=\"M42,-88C42,-88 12,-88 12,-88 6,-88 0,-82 0,-76 0,-76 0,-64 0,-64 0,-58 6,-52 12,-52 12,-52 42,-52 42,-52 48,-52 54,-58 54,-64 54,-64 54,-76 54,-76 54,-82 48,-88 42,-88\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-65.8\" font-family=\"Times,serif\" font-size=\"14.00\">input</text>\n",
       "</g>\n",
       "<!-- input&#45;&gt;prep/conv -->\n",
       "<g id=\"edge1\" class=\"edge\"><title>input&#45;&gt;prep/conv</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M54.4029,-70C62.3932,-70 71.3106,-70 79.8241,-70\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"79.919,-73.5001 89.919,-70 79.919,-66.5001 79.919,-73.5001\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<core.DotGraph at 0x10ea01710>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "colors = ColorMap()\n",
    "draw = lambda graph: DotGraph({p: ({'fillcolor': colors[type(v)], 'tooltip': repr(v)}, inputs) for p, (v, inputs) in graph.items() if v is not None})\n",
    "\n",
    "draw(build_graph(net()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "NB: on the first run, the first epoch will be slower as initialisation and Cudnn benchmarking take place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Run 0 at 2018-12-04 16:36:29\n",
      "       epoch           lr   train time   train loss    train acc    test time    test loss     test acc   total time\n",
      "           1       0.0800       4.2188       1.6095       0.4167       0.4228       1.4103       0.5157       4.6416\n",
      "           2       0.1600       2.9172       0.9291       0.6712       0.2854       0.8575       0.7107       7.8441\n",
      "           3       0.2400       2.9170       0.7249       0.7455       0.2260       0.8204       0.7263      10.9871\n",
      "           4       0.3200       2.9177       0.6237       0.7836       0.2259       0.6574       0.7718      14.1306\n",
      "           5       0.4000       2.9189       0.5514       0.8089       0.2262       0.7334       0.7503      17.2757\n",
      "           6       0.3789       2.9205       0.4934       0.8316       0.2262       0.6091       0.7845      20.4224\n",
      "           7       0.3579       2.9619       0.4416       0.8492       0.2261       0.4791       0.8369      23.6104\n",
      "           8       0.3368       2.9238       0.4107       0.8600       0.2264       0.5243       0.8217      26.7606\n",
      "           9       0.3158       2.9209       0.3874       0.8679       0.2258       0.5517       0.8121      29.9073\n",
      "          10       0.2947       2.9241       0.3599       0.8786       0.2261       0.4729       0.8363      33.0575\n",
      "          11       0.2737       2.9642       0.3402       0.8858       0.2259       0.4199       0.8644      36.2476\n",
      "          12       0.2526       2.9216       0.3243       0.8899       0.2266       0.5120       0.8322      39.3959\n",
      "          13       0.2316       2.9218       0.3066       0.8973       0.2261       0.4447       0.8503      42.5438\n",
      "          14       0.2105       2.9224       0.2883       0.9030       0.2267       0.3432       0.8850      45.6928\n",
      "          15       0.1895       2.9679       0.2662       0.9104       0.2268       0.3855       0.8736      48.8876\n",
      "          16       0.1684       2.9227       0.2490       0.9164       0.2265       0.6563       0.7905      52.0367\n",
      "          17       0.1474       2.9212       0.2278       0.9234       0.2261       0.3008       0.8965      55.1841\n",
      "          18       0.1263       2.9252       0.2047       0.9314       0.2259       0.2919       0.9010      58.3352\n",
      "          19       0.1053       2.9607       0.1871       0.9369       0.2257       0.2979       0.9014      61.5216\n",
      "          20       0.0842       2.9148       0.1597       0.9469       0.2262       0.2387       0.9184      64.6626\n",
      "          21       0.0632       2.9163       0.1373       0.9551       0.2258       0.2231       0.9255      67.8047\n",
      "          22       0.0421       2.9182       0.1139       0.9648       0.2261       0.2046       0.9331      70.9490\n",
      "          23       0.0211       2.9512       0.0885       0.9736       0.2258       0.1845       0.9393      74.1260\n",
      "          24       0.0000       2.9243       0.0751       0.9783       0.2257       0.1741       0.9437      77.2760\n",
      "Starting Run 1 at 2018-12-04 16:37:50\n",
      "       epoch           lr   train time   train loss    train acc    test time    test loss     test acc   total time\n",
      "           1       0.0800       2.9244       1.6165       0.4156       0.2266       1.2236       0.5618       3.1510\n",
      "           2       0.1600       2.9343       0.9368       0.6655       0.2265       1.2664       0.5835       6.3118\n",
      "           3       0.2400       2.9752       0.7302       0.7455       0.2269       0.8205       0.7104       9.5139\n",
      "           4       0.3200       2.9295       0.6205       0.7851       0.2280       0.6482       0.7829      12.6714\n",
      "           5       0.4000       2.9263       0.5544       0.8083       0.2276       0.6061       0.7916      15.8253\n",
      "           6       0.3789       2.9283       0.4978       0.8295       0.2270       0.5547       0.8149      18.9806\n",
      "           7       0.3579       2.9683       0.4426       0.8494       0.2265       0.5598       0.8020      22.1753\n",
      "           8       0.3368       2.9276       0.4075       0.8605       0.2269       0.4830       0.8346      25.3299\n",
      "           9       0.3158       2.9273       0.3828       0.8695       0.2265       0.4121       0.8591      28.4837\n",
      "          10       0.2947       2.9257       0.3603       0.8768       0.2263       0.5355       0.8191      31.6357\n",
      "          11       0.2737       2.9684       0.3494       0.8806       0.2266       0.4747       0.8375      34.8307\n",
      "          12       0.2526       2.9248       0.3239       0.8903       0.2274       0.4424       0.8542      37.9829\n",
      "          13       0.2316       2.9226       0.3079       0.8953       0.2274       0.3844       0.8717      41.1329\n",
      "          14       0.2105       2.9245       0.2867       0.9021       0.2263       0.3982       0.8660      44.2836\n",
      "          15       0.1895       2.9560       0.2717       0.9087       0.2267       0.4928       0.8381      47.4663\n",
      "          16       0.1684       2.9256       0.2475       0.9158       0.2265       0.3210       0.8914      50.6183\n",
      "          17       0.1474       2.9261       0.2273       0.9230       0.2260       0.3258       0.8934      53.7705\n",
      "          18       0.1263       2.9252       0.2082       0.9306       0.2262       0.4067       0.8660      56.9219\n",
      "          19       0.1053       2.9661       0.1854       0.9386       0.2260       0.2674       0.9117      60.1140\n",
      "          20       0.0842       2.9241       0.1675       0.9455       0.2260       0.2542       0.9161      63.2642\n",
      "          21       0.0632       2.9258       0.1413       0.9547       0.2263       0.2174       0.9269      66.4163\n",
      "          22       0.0421       2.9263       0.1142       0.9638       0.2262       0.2254       0.9268      69.5688\n",
      "          23       0.0211       2.9659       0.0930       0.9716       0.2262       0.1968       0.9359      72.7609\n",
      "          24       0.0000       2.9224       0.0745       0.9795       0.2265       0.1745       0.9433      75.9098\n",
      "Starting Run 2 at 2018-12-04 16:39:06\n",
      "       epoch           lr   train time   train loss    train acc    test time    test loss     test acc   total time\n",
      "           1       0.0800       2.9232       1.6216       0.4130       0.2269       1.1244       0.5880       3.1501\n",
      "           2       0.1600       2.9312       0.9384       0.6649       0.2269       0.9665       0.6584       6.3082\n",
      "           3       0.2400       2.9759       0.7258       0.7470       0.2267       0.7609       0.7427       9.5107\n",
      "           4       0.3200       2.9323       0.6152       0.7852       0.2272       0.8803       0.6955      12.6702\n",
      "           5       0.4000       2.9296       0.5548       0.8098       0.2269       0.5719       0.8027      15.8267\n",
      "           6       0.3789       2.9264       0.4938       0.8296       0.2272       0.5524       0.8045      18.9803\n",
      "           7       0.3579       2.9669       0.4409       0.8482       0.2268       0.4810       0.8354      22.1741\n",
      "           8       0.3368       2.9273       0.4106       0.8604       0.2267       0.4488       0.8440      25.3280\n",
      "           9       0.3158       2.9282       0.3889       0.8682       0.2260       0.5664       0.8090      28.4823\n",
      "          10       0.2947       2.9231       0.3619       0.8765       0.2266       0.4685       0.8406      31.6320\n",
      "          11       0.2737       2.9686       0.3440       0.8830       0.2256       0.4093       0.8576      34.8262\n",
      "          12       0.2526       2.9237       0.3255       0.8897       0.2266       0.3533       0.8802      37.9764\n",
      "          13       0.2316       2.9223       0.3075       0.8944       0.2269       0.3974       0.8631      41.1256\n",
      "          14       0.2105       2.9225       0.2880       0.9023       0.2267       0.5371       0.8259      44.2747\n",
      "          15       0.1895       2.9665       0.2664       0.9101       0.2262       0.4167       0.8565      47.4674\n",
      "          16       0.1684       2.9245       0.2534       0.9153       0.2266       0.3200       0.8920      50.6184\n",
      "          17       0.1474       2.9247       0.2301       0.9215       0.2263       0.3205       0.8919      53.7695\n",
      "          18       0.1263       2.9228       0.2082       0.9304       0.2264       0.2819       0.9065      56.9187\n",
      "          19       0.1053       2.9697       0.1859       0.9383       0.2261       0.3093       0.9001      60.1145\n",
      "          20       0.0842       2.9239       0.1612       0.9475       0.2260       0.2465       0.9180      63.2644\n",
      "          21       0.0632       2.9242       0.1407       0.9542       0.2262       0.2437       0.9192      66.4148\n",
      "          22       0.0421       2.9218       0.1168       0.9632       0.2266       0.2360       0.9219      69.5632\n",
      "          23       0.0211       2.9640       0.0922       0.9727       0.2259       0.1939       0.9352      72.7531\n",
      "          24       0.0000       2.9200       0.0768       0.9783       0.2264       0.1758       0.9404      75.8995\n",
      "Starting Run 3 at 2018-12-04 16:40:22\n",
      "       epoch           lr   train time   train loss    train acc    test time    test loss     test acc   total time\n",
      "           1       0.0800       2.9281       1.6001       0.4234       0.2281       1.4474       0.5124       3.1561\n",
      "           2       0.1600       2.9342       0.9118       0.6754       0.2266       0.9700       0.6833       6.3169\n",
      "           3       0.2400       2.9822       0.7220       0.7478       0.2274       0.9374       0.6959       9.5264\n",
      "           4       0.3200       2.9353       0.6156       0.7847       0.2271       0.7660       0.7489      12.6889\n",
      "           5       0.4000       2.9306       0.5565       0.8061       0.2272       0.6762       0.7650      15.8467\n",
      "           6       0.3789       2.9255       0.4974       0.8274       0.2271       0.5478       0.8077      18.9992\n",
      "           7       0.3579       2.9686       0.4384       0.8490       0.2271       0.7755       0.7329      22.1949\n",
      "           8       0.3368       2.9291       0.4096       0.8588       0.2270       0.7280       0.7418      25.3510\n",
      "           9       0.3158       2.9292       0.3787       0.8696       0.2268       0.4444       0.8510      28.5071\n",
      "          10       0.2947       2.9282       0.3638       0.8754       0.2268       0.4216       0.8563      31.6621\n",
      "          11       0.2737       2.9578       0.3388       0.8847       0.2262       0.4204       0.8547      34.8461\n",
      "          12       0.2526       2.9290       0.3243       0.8885       0.2265       0.3991       0.8680      38.0016\n",
      "          13       0.2316       2.9269       0.3015       0.8991       0.2267       0.4715       0.8509      41.1552\n",
      "          14       0.2105       2.9294       0.2867       0.9022       0.2264       0.3760       0.8754      44.3110\n",
      "          15       0.1895       2.9494       0.2687       0.9084       0.2266       0.4472       0.8562      47.4870\n",
      "          16       0.1684       2.9263       0.2462       0.9179       0.2263       0.5258       0.8224      50.6395\n",
      "          17       0.1474       2.9238       0.2335       0.9208       0.2265       0.3914       0.8718      53.7899\n",
      "          18       0.1263       2.9281       0.2088       0.9294       0.2265       0.3109       0.8933      56.9444\n",
      "          19       0.1053       2.9483       0.1876       0.9367       0.2270       0.3537       0.8839      60.1197\n",
      "          20       0.0842       2.9237       0.1671       0.9446       0.2265       0.2555       0.9186      63.2699\n",
      "          21       0.0632       2.9249       0.1402       0.9551       0.2266       0.2182       0.9266      66.4213\n",
      "          22       0.0421       2.9271       0.1122       0.9641       0.2262       0.1989       0.9357      69.5747\n",
      "          23       0.0211       2.9650       0.0913       0.9727       0.2265       0.1835       0.9399      72.7662\n",
      "          24       0.0000       2.9277       0.0739       0.9789       0.2264       0.1729       0.9423      75.9203\n",
      "Starting Run 4 at 2018-12-04 16:41:38\n",
      "       epoch           lr   train time   train loss    train acc    test time    test loss     test acc   total time\n",
      "           1       0.0800       2.9241       1.6148       0.4175       0.2274       1.6599       0.4501       3.1515\n",
      "           2       0.1600       2.9326       0.9273       0.6689       0.2268       0.8414       0.7104       6.3109\n",
      "           3       0.2400       2.9629       0.7208       0.7477       0.2277       1.0110       0.6613       9.5015\n",
      "           4       0.3200       2.9303       0.6254       0.7828       0.2273       0.7714       0.7402      12.6591\n",
      "           5       0.4000       2.9315       0.5565       0.8069       0.2276       0.6201       0.7875      15.8182\n",
      "           6       0.3789       2.9284       0.4928       0.8287       0.2273       0.5163       0.8195      18.9740\n",
      "           7       0.3579       2.9693       0.4435       0.8480       0.2270       0.4975       0.8339      22.1702\n"
     ]
    }
   ],
   "source": [
    "epochs=24\n",
    "lr_schedule = PiecewiseLinear([0, 5, epochs], [0, 0.4, 0])\n",
    "batch_size = 512\n",
    "train_transforms = [Crop(32, 32), FlipLR(), Cutout(8, 8)]\n",
    "N_runs = 3\n",
    "\n",
    "train_batches = DataLoader(Transform(train_set, train_transforms), batch_size, shuffle=True, set_random_choices=True, drop_last=True)\n",
    "valid_batches = DataLoader(valid_set, batch_size, shuffle=False, drop_last=False)\n",
    "lr = lambda step: lr_schedule(step/len(train_batches))/batch_size\n",
    "\n",
    "summaries = []\n",
    "for i in range(N_runs):\n",
    "    print(f'Starting Run {i} at {localtime()}')\n",
    "    model = Network(net()).to(device).half()\n",
    "    opts = [SGD(trainable_params(model).values(), {'lr': lr, 'weight_decay': Const(5e-4*batch_size), 'momentum': Const(0.9)})]\n",
    "    logs, state = Table(), {MODEL: model, LOSS: x_ent_loss, OPTS: opts}\n",
    "    for epoch in range(epochs):\n",
    "        logs.append(union({'epoch': epoch+1}, train_epoch(state, Timer(torch.cuda.synchronize), train_batches, valid_batches)))\n",
    "logs.df().query(f'epoch=={epochs}')[['train_acc', 'valid_acc']].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
